
from torch.utils.data import Dataset, DataLoader

import json
import math
from random import random
import torch
import sys 
sys.path.append('..')
from util.graph import Graph
class GraphDataset(Dataset):
    """
    数据集类，用于加载和预处理 Graph 对象的 JSON 文件。
    """
    def __init__(self, graph_record_path,split_path,split_name):
        """
        初始化数据集。
        :param json_paths: 包含 Graph JSON 文件路径的列表。
        :param max_length: 用于归一化边长的最大长度。
        """
        with open(graph_record_path, 'r') as f:
            self.graph_records = json.load(f)
            
        with open(split_path, 'r') as f:
            self.split_list = json.load(f)[split_name]
    def build_gt(self,node_map):
        """
        构建图的 ground truth。
        :param node_map: 节点映射字典。
        :return: ground truth 矩阵。
        """
        num_node_1 = len(node_map)
        num_node_2 = max(node_map.keys()) + 1
        gt = torch.zeros((num_node_1, num_node_2))
        for node_id, node in node_map.items():
            for neighbor_id in node.neighbors:
                gt[node_id, neighbor_id] = 1
        return gt
    def __len__(self):
        return len(self.graphs)

    def __getitem__(self, idx):
        """
        返回图数据。
        :param idx: 数据索引。
        :return: 图节点、边和相关属性。
        """
        image_name = self.split_list[idx]
        graph_record=self.graph_records[image_name]
        transformed_number=graph_record['transform_number']
        selected_index=int(random()*transformed_number+1)
        select_json=graph_record['data'][selected_index]
        graph0 = Graph().load(graph_record['data'][0])
        graph1 = Graph().load(select_json)
        if 'node_map' in graph_record:
            node_map = graph_record['node_map'][selected_index]
            gt = self.build_gt(node_map)
        else:
            raise ValueError('No node map found in graph record.')
        return graph0, graph1, gt